from torch import cuda
import os
#os.environ['DATASET_DIR'] = "C:/pytorch training/NEU1"
#os.environ['DATASET_DIR'] = "C:/pytorch training/omniglot_dataset"
#os.environ['DATASET_DIR'] = "C:/pytorch training/mini-imagenet-1"

def get_args():
    import argparse
    import torch

    parser = argparse.ArgumentParser(description='Meta-Learning Training System')

    #################################### 基础配置 ####################################
    parser.add_argument('--experiment_name', type=str, default="MeTal", help='实验名称')
    parser.add_argument('--continue_from_epoch', type=str, default='from_scratch',
                       choices=['from_scratch', 'latest', 'specific_epoch'], help='恢复训练模式')
    parser.add_argument('--name_of_args_json_file', type=str, default="None", help='JSON配置文件路径')
    parser.add_argument('--architecture_name', type=str, help='模型架构名称')

    #################################### 设备配置 ####################################
    parser.add_argument('--gpu_to_use', type=int, default=0, help='指定使用的GPU编号')
    parser.add_argument('--num_of_gpus', type=int, default=1, help='使用的GPU数量')

    #################################### 数据配置 ####################################
    #"mini-imagenet""C:/pytorch training/mini-imagenet-1""NEU""C:/pytorch training/NEU1"
    parser.add_argument('--train_dataset_name', type=str, default="NEU", help='训练集名称')
    parser.add_argument('--train_dataset_path', type=str, default="C:/pytorch training/NEU1",
                       help='训练集路径（自动分割训练/验证集）')
    parser.add_argument('--test_dataset_name', type=str, default="NEU", help='测试集名称')
    parser.add_argument('--test_dataset_path', type=str, default="C:/pytorch training/NEU1",
                       help='独立测试集路径')
    parser.add_argument('--train_val_split_ratio', type=float, default=0.9, help='训练/验证集分割比例')
    parser.add_argument('--image_size', type=int, nargs=2, default=[224, 224], help='输入图像尺寸[高, 宽]')
    parser.add_argument('--image_channels', type=int, default=1, choices=[1, 3], help='输入通道数')
    parser.add_argument('--indexes_of_folders_indicating_class', type=int, nargs='+', default=[-2, -3],
                       help='类别文件夹路径索引（例如[-2,-3]对应路径/dataset_root/images/class/img.jpg）')

    #################################### 模型架构 ####################################
    parser.add_argument('--backbone', type=str, default='4-CONV', choices=['4-CONV', 'ResNet12'], help='骨干网络架构')
    parser.add_argument('--num_filters', type=int, default=64, help='CNN每层滤波器数量')
    parser.add_argument('--norm_layer', type=str, default="batch_norm",
                       choices=['batch_norm', 'layer_norm', 'none'], help='归一化层类型')
    parser.add_argument('--max_pooling', type=lambda x: x.lower() == "true", default=False, help='是否使用最大池化')
    parser.add_argument('--conv_padding', type=int, default=1, help='卷积层填充大小')
    parser.add_argument('--num_stages', type=int, default=3, help='网络阶段数')
    parser.add_argument('--cnn_num_blocks', type=int, default=4, help='CNN块总数')
    parser.add_argument('--cnn_blocks_per_stage', type=int, default=1, help='每个阶段的CNN块数')
    parser.add_argument('--dropout_rate_value', type=float, default=0.3, help='Dropout概率')

    #################################### 训练参数 ####################################
    parser.add_argument('--batch_size', type=int, default=32, help='全局批次大小')
    parser.add_argument('--num_epochs', type=int, default=100, help='总训练轮数')
    parser.add_argument('--total_iter_per_epoch', type=int, default=500, help='每轮迭代次数')
    parser.add_argument('--samples_per_iter', type=int, default=1, help='每次迭代样本数')
    parser.add_argument('--meta_learning_rate', type=float, default=0.001, help='元学习率')
    parser.add_argument('--task_learning_rate', type=float, default=0.01, help='内部任务学习率')
    parser.add_argument('--min_learning_rate', type=float, default=0.00001, help='最小学习率')
    parser.add_argument('--number_of_training_steps_per_iter', type=int, default=5,
                       help='每次迭代的内部训练步数')
    parser.add_argument('--number_of_evaluation_steps_per_iter', type=int, default=1,
                       help='每次迭代的内部评估步数')
    parser.add_argument('--total_epochs_before_pause', type=int, default=200, help='训练暂停前的总轮数')

    #################################### 优化器参数 ####################################
    parser.add_argument('--init_inner_loop_learning_rate', type=float, default=0.001, help='内部循环初始学习率')
    parser.add_argument('--init_inner_loop_weight_decay', type=float, default=0.01, help='内部循环权重衰减')
    parser.add_argument('--learnable_per_layer_per_step_inner_loop_learning_rate',
                       action='store_true', help='分层自适应学习率')
    parser.add_argument('--use_learnable_learning_rates', type=lambda x: x.lower() == "true",
                       default=True, help='是否使用可学习学习率')

    #################################### Few-shot 设置 ####################################
    parser.add_argument('--num_ways', type=int, default=5, help='每个任务类别数(n-way)')
    parser.add_argument('--num_shots', type=int, default=5, help='每类支持样本数(k-shot)')
    parser.add_argument('--num_query', type=int, default=5, help='每类查询样本数')
    parser.add_argument('--num_evaluation_tasks', type=int, default=128, help='评估任务数量')

    #################################### 高级训练选项 ####################################
    parser.add_argument('--use_multi_step_loss_optimization', type=lambda x: x.lower() == "true",
                       default=False, help='是否使用多步损失优化')
    parser.add_argument('--multi_step_loss_num_epochs', type=int, default=20,
                       help='多步损失优化的轮数')
    parser.add_argument('--second_order',
                        type=lambda x: x.lower() == "true",
                        default=False,
                        help='是否使用二阶优化（MAML核心参数）')

    #################################### 系统配置 ####################################
    parser.add_argument('--num_workers', type=int, default=4, help='数据加载线程数')
    parser.add_argument('--max_models_to_save', type=int, default=5, help='最大保存模型数')

    #################################### 随机种子 ####################################
    parser.add_argument('--seed', type=int, default=104, help='全局随机种子')
    parser.add_argument('--train_seed', type=int, default=42, help='训练集随机种子')
    parser.add_argument('--val_seed', type=int, default=42, help='验证集随机种子')
    parser.add_argument('--test_seed', type=int, default=42, help='测试集随机种子')

    #################################### BN 参数 ####################################
    parser.add_argument('--learnable_bn_gamma', action='store_true', help='可学习BN gamma参数')
    parser.add_argument('--learnable_bn_beta', action='store_true', help='可学习BN beta参数')
    parser.add_argument('--enable_inner_loop_optimizable_bn_params', action='store_true',
                       help='内部循环优化BN参数')
    parser.add_argument('--per_step_bn_statistics', type=lambda x: x.lower() == "true",
                       default=False, help='是否启用每步BN统计')

    #################################### 布尔参数修复 ####################################
    bool_params = [
        ('load_into_memory', 'False'),  # 关键修复：添加缺失的参数
        ('reset_stored_filepaths', 'True'),
        ('reset_stored_paths', 'False'),
        ('reverse_channels', 'False'),
        ('labels_as_int', 'False'),
        ('evaluate_on_test_set_only', 'False'),
        ('meta_opt_bn', 'False'),
        ('attenuate', 'False'),
        ('alfa', 'False'),
        ('random_init', 'False'),
        ('meta_loss', 'True'),
    ]
    for arg, default in bool_params:
        parser.add_argument(f'--{arg}',
                           type=lambda x: x.lower() == "true",
                           default=default.lower() == "true",
                           help=f'布尔参数: {arg} (默认: {default})')

    #################################### 参数后处理 ####################################
    args = parser.parse_args()

    # 自动生成参数
    args.image_height, args.image_width = args.image_size
    args.num_classes_per_set = args.num_ways
    args.num_samples_per_class = args.num_shots
    args.num_target_samples = args.num_query
    args.num_dataprovider_workers = args.num_workers
    args.total_epochs = args.num_epochs
    args.cnn_num_filters = args.num_filters

    # 设备配置
    args.use_cuda = torch.cuda.is_available()
    args.device = torch.device(f'cuda:{args.gpu_to_use}' if args.use_cuda else 'cpu')

    return args, args.device



class Bunch(object):
  def __init__(self, adict):
    self.__dict__.update(adict)

def extract_args_from_json(json_file_path, args_dict):
    import json
    summary_filename = json_file_path
    with open(summary_filename) as f:
        summary_dict = json.load(fp=f)

    for key in summary_dict.keys():
        if "continue_from" in key:
            pass
        elif "gpu_to_use" in key:
            pass
        else:
            args_dict[key] = summary_dict[key]

    return args_dict





